from marker.config.parser import ConfigParser
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered

from markitdown_api.app.handler.task_handler import TaskHandler, TaskResult
from markitdown_api.app.models.task import Task


class MarkerPdfHandler(TaskHandler):
    """
    基于 marker pdf 处理 pdf
    @see https://github.com/vikParuchuri/marker
    """

    def __init__(self):
        # 模型配置
        self.model_dict = create_model_dict()

        config = {
            "output_format": "markdown",
            # 不要从 PDF 中提取图片
            "disable_image_extraction": True
        }
        self.config_parser = ConfigParser(config)


def handle(self, task: Task) -> TaskResult:
    converter = PdfConverter(
        config=self.config_parser.generate_config_dict(),
        artifact_dict=self.model_dict,
        processor_list=self.config_parser.get_processors(),
        renderer=self.config_parser.get_renderer(),
        llm_service=self.config_parser.get_llm_service(),
    )
    rendered = converter(task.input_path)
    text, _, _ = text_from_rendered(rendered)
    return TaskResult(
        content=text
    )
